structure UNIFICATION = struct
structure A = ABSTRACT
open UTIL

(* =============================================================================================
 * Unification of type variables
 * ============================================================================================= *)
exception UNIFY;

(* app: (int * ty) list -> int -> ty *)
fun app ((z,t)::s) x = if x = z then t else app s x
  | app [] x         = raise TypeError("Match failure in TYPECHECK.app on args: [] " ^ Int.toString x)

(* lift: (int * ty) list -> ty -> ty *)
fun lift s = mapTy (fn (t as TyTerm(r,[TyVar x])) => if List.exists (fn (i,_) => x = i) s
                                                     then let
                                                              val newTy = app s x
                                                          in
                                                              case newTy of
                                                                TyTerm _ => newTy
                                                              | baseTy => TyTerm(r,[baseTy])
                                                          end
                                                     else t
                     | otherTy   => otherTy )

(* solve: (ty * ty) list * env -> env *)
fun solve [] e = e
  | solve ((TyVar var,  ty2       )::xs) e = elim [(var,ty2)] xs e
  | solve (( ty1,       TyVar var )::xs) e = elim [(var,ty1)] xs e
  | solve ((TySum(a,b), TySum(c,d))::xs) e = solve ((a,c)::(b,d)::xs) e
  | solve ((TyRule(a,b),TyRule(c,d))::xs)e = solve ((a,c)::(b,d)::xs) e
  | solve ((a as TyTerm (x,xs),
            b as TyTerm (y,ys))::rest) e
      = ((*fine("unifying " ^ toString a ^ " with " ^ toString b) 2;*)
        if x <> y andalso
           x <> "'term" andalso y <> "'term" (* neither term's root is generic *)
        then raise UNIFY
        else let
                 fun zip [] [] = []
                   | zip [] _  = raise UNIFY
                   | zip _  [] = raise UNIFY
                   | zip (b::bs) (c::cs) = (b,c)::(zip bs cs)
                 fun getFirstMatch (ty1::rest) tySome
                       = let
                             val (unified, env) = (true, solve [(ty1,tySome)] e) handle UNIFY => (false, e)
                         in
                             if unified
                             then env
                             else getFirstMatch rest tySome
                         end
                   | getFirstMatch [] _ = raise UNIFY
                 fun existsChoice kids (TySum(a,b)) = (existsChoice kids a handle UNIFY => existsChoice kids b)
                   | existsChoice kids anyOther = getFirstMatch kids anyOther
             in
                 case xs of 
                   [TyVar var] => (case ys of
                                     [TyVar z]  => (*let val newVar = nextVar()
                                                   in  elim [(var,newVar), (z,newVar)] rest e
                                                   end*)
                                                   elim [(z,TyVar var)] rest e
                                   | [TySome s] => elim [(var,TySome s)] rest e
                                   | _          => elim [(var,b)] rest e)
                 | [TySome s] => (case ys of
                                    [TyVar var] => elim [(var,TySome s)] rest e
                                  | [TySome s2] => solve ((s,s2)::rest) e
                                  | _ => if List.exists (fn i => i = s) ys
                                         then solve rest e
                                         else existsChoice ys s)
                 | _ => case ys of
                          [TyVar var] => elim [(var,a)] rest e
                        | [TySome s] => if List.exists (fn i => i = s) xs
                                        then solve rest e
                                        else existsChoice xs s
                 | _ => solve ((zip xs ys) @ rest) e
             end)
  | solve ((TyTerm _,        TySum  _        )::xs) e = raise UNIFY
  | solve ((TySum  _,        TyTerm _        )::xs) e = raise UNIFY
  | solve ((TyTerm _,        TyRule _        )::xs) e = raise UNIFY
  | solve ((TyRule _,        TyTerm _        )::xs) e = raise UNIFY
  | solve ((TyTerm(_,[ty1]), ty2             )::xs) e = solve ((ty1,ty2)::xs) e (* linear term and a base type *)
  | solve (( ty1,            TyTerm (_,[ty2]))::xs) e = solve ((ty1,ty2)::xs) e

  | solve ((a as TyTerm(x,[]),ty2            )::xs) e = (case inferPrimType a of
                                                           TyTerm(x,[]) => if a = ty2 then solve xs e
                                                                           else raise UNIFY
                                                         | z => solve ((z, ty2)::xs) e) (* leaf and prim type *)
  | solve ((ty1,             TyTerm(y,[])    )::xs) e = solve ((TyTerm(y,[]), ty1)::xs) e

  | solve ((r as TyRule _,   TySum(a,b)      )::xs) e = (solve ((r,a)::xs) e handle
                                                          UNIFY => solve ((r,b)::xs) e)

  | solve (( ty1,            ty2             )::xs) e = if ty1 = ty2
                                                        then solve xs e
                                                        else raise UNIFY

(* elim: (int * ty) list -> (ty * ty) list -> env -> env *)
and elim []          tts e = solve tts e
  | elim ((x,t)::xs) tts e = let
                                 val xt = lift [(x,t)]
                                 val tts2 = map (fn (tyA,tyB) => (xt tyA, xt tyB)) tts
                                 val e2 = S.insert(e, x, t)
                                 val oldBindings : (int * ty) list = S.listItemsi e2
                                 (*fun update ((key,TyVar oldVar)::rest) e = if oldVar = x
                                                                           then update rest (S.insert(e,key,t))
                                                                           else update rest e
                                   | update ( _                ::rest) e = update rest e
                                   | update [] e = e*)
                                 fun update ((key,ty1)::rest) e = update rest (S.insert(e,key,applySubst e ty1))
                                   | update [] e = e
                                 val e3 = update oldBindings e2
                             in
                                 elim xs tts2 e3
                             end

(* Unifies a type with a list of expected types. The first successful unification returns an updated context.
 * unify: (ty * ty list * context * expr) -> context *)
fun unify (givenTy, [expectedTy], env, expr)
      = ((solve [(givenTy,expectedTy)] env) handle
                     UNIFY => raiseOperandError (givenTy,expectedTy,expr))
  | unify (givenTy, expectedTypes, env, expr)
      = let
            fun polyUnify  []           = raisePolyOpError (givenTy, expr)
              | polyUnify (expTy::rest) = (solve [(givenTy,expTy)] env) handle UNIFY => polyUnify rest

        in
            polyUnify expectedTypes
        end
end